import cv2
import numpy as np

class ColorDetector(object):
    """颜色识别器
    * 通过调节trackbar来调节颜色阈值
    * 使用detect方法来识别颜色，得到二值化过滤后的图像数据
    * 使用draw_rectangle方法来在图像上绘制矩形框，传入的应该是二值化过滤后的图像数据"""
    def __init__(self) -> None:
        """初始化"""
        # region 初始颜色阈值
        self.low_h = 0
        self.low_s = 0
        self.low_v = 0

        self.high_h = 180
        self.high_s = 255
        self.high_v = 255

        self.minarea = 0
        self.maxarea = 100000 # 220800
        # endregion

    def __call__(self,id:int=0):
        """使用call方法调出trackbar
        * id: 识别器的id的id"""
        # region 创建trackbar
        cv2.namedWindow(f'trackbar{id}', cv2.WINDOW_NORMAL)
        cv2.createTrackbar('low_h', f'trackbar{id}', self.low_h, 180, self.call_back_low_h)
        cv2.createTrackbar('high_h', f'trackbar{id}', self.high_h, 180, self.call_back_high_h)
        cv2.createTrackbar('low_s', f'trackbar{id}', self.low_s, 255, self.call_back_low_s)
        cv2.createTrackbar('high_s', f'trackbar{id}', self.high_s, 255, self.call_back_high_s)
        cv2.createTrackbar('low_v', f'trackbar{id}', self.low_v, 255, self.call_back_low_v)
        cv2.createTrackbar('high_v', f'trackbar{id}', self.high_v, 255, self.call_back_high_v)
        cv2.createTrackbar('minarea', f'trackbar{id}', self.minarea, 100000, self.call_back_minarea)
        cv2.createTrackbar('maxarea', f'trackbar{id}', self.maxarea, 100000, self.call_back_maxarea)
        # endregion
        pass

    # region trackbar回调函数
    def callback(self, x):
        pass

    def call_back_low_h(self, low_h):
        self.low_h = low_h
    
    def call_back_high_h(self, high_h):
        self.high_h = high_h

    def call_back_low_s(self, low_s):
        self.low_s = low_s

    def call_back_high_s(self, high_s):
        self.high_s = high_s

    def call_back_low_v(self, low_v):
        self.low_v = low_v

    def call_back_high_v(self, high_v):
        self.high_v = high_v

    def call_back_minarea(self, minarea):
        self.minarea = minarea

    def call_back_maxarea(self, maxarea):
        self.maxarea = maxarea
    # endregion

    def detect(self, img:cv2.typing.MatLike):
        """颜色识别
        * img: 传入的图像数据
        返回值：二值化过滤后的图像数据"""
        _shape = img.shape
        try:
            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)				# 将BGR图像转换成HSV图像
        except:
            return np.zeros((_shape[0], _shape[1]), np.uint8)

        low = np.array([self.low_h, self.low_s, self.low_v])		# 低阈值
        high = np.array([self.high_h, self.high_s, self.high_v])	# 高阈值

        mask = cv2.inRange(hsv, low, high)						# 通过阈值过滤图像，将在阈值范围内的像素点设置为255，不在阈值范围内的像素点设置为0
        kernel = np.ones((5, 5), np.uint8)						# 创建一个5*5的矩阵，矩阵元素全为1
        opencal = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)	# 开运算，先腐蚀后膨胀
        # 对opencal进行腐蚀操作，去除噪声
        res = cv2.erode(opencal, kernel, iterations=1)
        return res

    def draw_rectangle(self, img:cv2.typing.MatLike): 
        """在图像上绘制矩形框
        * img: 传入的二值化图像数据
        返回值：绘制矩形框后的图像数据，矩形框的坐标"""
        lst = []
        for cnt in ColorDetector._get_edge(img):										# 遍历轮廓数据
            x, y, w, h = cv2.boundingRect(cnt)						# 获取矩形框的坐标和宽高

            if (self.minarea > w*h or w*h > self.maxarea):
                continue
            if h/w > 1.5:
                continue
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)	# 在图像上绘制矩形框
            lst.append((x, y, x + w, y + h))
        return img, lst
    
    @staticmethod
    def _get_edge(img:cv2.typing.MatLike):
        """静态函数获取边缘
        * img: 传入的图像数据
        返回值：边缘图像"""
        # conuours是轮廓，hierarchy是轮廓的层级
        contours, _ = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        return contours
    
    def __del__(self):
        """析构函数"""
        cv2.destroyAllWindows()
        pass

class RectangleDetector(object):
    """多边形识别器"""
    def __init__(self) -> None:
        self.Binarization = 113  #二值化参数
        self.canny_min = 0
        self.canny_max = 0
        self.minarea = 0                # 最小面积
        self.maxarea = 100000           # 最大面积
        self.LWratio = None             # 长宽比默认不设置
        self.ori_img = None             # 原图像
        pass

    def __call__(self,id:int=0):
        """使用call方法调出trackbar"""

        # region 创建trackbar
        cv2.namedWindow(f'trackbar{id}', cv2.WINDOW_NORMAL)
        cv2.createTrackbar('minarea', f'trackbar{id}', self.minarea, 100000, self.call_back_minarea)
        cv2.createTrackbar('maxarea', f'trackbar{id}', self.maxarea, 100000, self.call_back_maxarea)
        # cv2.createTrackbar('LWratio', f'trackbar{id}', 0, 100, self.call_back_LWratio)

        pass

    # region trackbar回调函数
    def callback(self, x):
        pass

    def call_back_minarea(self, minarea):
        self.minarea = minarea

    def call_back_maxarea(self, maxarea):
        self.maxarea = maxarea

    def call_back_LWratio(self, LWratio):
        self.LWratio = LWratio
    # endregion
        
    def _extract_the_edges(self, img:cv2.typing.MatLike):
        """提取边缘
        * img: 原图像
        返回值：边缘图像"""
        self.ori_img = img

        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        img = cv2.GaussianBlur(img, (5, 5), sigmaX=0.6)             #高斯滤波

        ret, img = cv2.threshold(img, self.Binarization, 255, cv2.THRESH_BINARY)  #二值化
        # cv2.imshow("2zh", img)

        img = cv2.Canny(img, self.canny_min, self.canny_max)                       #canny算子
        # cv2.imshow("canny", img)
        return img
    
    def find_polygon(self, img:cv2.typing.MatLike, _show=False):
        """多边形拟合
        * img: 二值化图像
        * _show: 是否显示图像,默认不显示
        * 返回值：多边形端点坐标，多边形的点数"""
        approx = None
        point_sum = 0           # 多边形的点数
        contours, _ = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)   #寻找轮廓

        for i in contours:
            area = cv2.contourArea(i)   #计算轮廓面积
            if self.minarea < area < self.maxarea:
                approx = cv2.approxPolyDP(i, 0.02 * cv2.arcLength(i, True), True)

                if self.ori_img is None:
                    raise ValueError("The original image is None,please use extract_the_edges method before find_rectangle.")
                
                cv2.drawContours(self.ori_img, [i], -1, (255, 255, 255), 2)
                # cv2.drawContours(img_new, [approx], -1, (255, 0, 0), 5)  # Fix: Pass a list of contours
                if _show:
                    cv2.imshow("img_new", self.ori_img)
                point_sum = point_sum + 1
        return approx, point_sum
    
    def find_rectangle(self, img:cv2.typing.MatLike, _show=False) -> np.ndarray|None:
        """矩形拟合
        ------
        * img: 图像
        * _show: 是否显示图像,默认不显示
        * 返回值：矩形的四个端点坐标"""
        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)     # 寻找轮廓
        for contour in contours:
            area = cv2.contourArea(contour)
            if self.minarea < area < self.maxarea:
                continue
            rect = cv2.minAreaRect(contour)
            box = cv2.boxPoints(rect)
            box = box.astype(int)

            if self.ori_img is None:
                raise ValueError("The original image is None,please use extract_the_edges method before find_rectangle.")
            
            cv2.drawContours(self.ori_img, [box], 0, (255, 0, 0), 5)        # 在原图上绘制矩形
            if _show:
                cv2.imshow("img_new", self.ori_img)
            return box

